#include "postgreV3_protocol.h"
#include "log.h"

namespace secsmart_protocol
{
    DETECT_RESULT PostgreProtocol::parse(secsmart_session::session_sptr session, const uint8_t *buf, size_t size, bool is_forward)
    {
        direction = secsmart_session::session_dirction_t::SESSION_DIRCTION_UNKNOW;
        if(size == 1)
        {
            int8_t m_cmdType = buf[0];
            switch(m_noCmdReqStatus)
            {
                case NoCmdReqStatus::SSLRequest_cmd:
                case NoCmdReqStatus::LoadBalance_cmd:
                {
                    if(m_noCmdReqStatus == NoCmdReqStatus::SSLRequest_cmd)
                    {
                        switch(m_cmdType)
                        {
                            case SSLRequestStatus::SSLError:
                            case SSLRequestStatus::SSLRefused:
                            case SSLRequestStatus::SSLOk:
                            {
                                LOG_DEBUG("SSLResponse packet detected !");
                                addMatchCount();
                                return DETECT_RESULT::DETECT_OK;
                            }
                            default:
                            {
                                return DETECT_RESULT::DETECT_UNKNOWN;
                            }
                        }
                    }
                    else
                    {
                        LOG_DEBUG("Other loadBalance Response packet detected !");
                        addMatchCount();
                        return DETECT_RESULT::DETECT_OK;
                    }
                    m_noCmdReqStatus = NoCmdReqStatus::Unknown_cmd;
                }
                default:
                {
                    m_noCmdReqStatus = NoCmdReqStatus::Unknown_cmd;
                    return DETECT_RESULT::DETECT_UNKNOWN;
                }
                
            }
        }
        int64_t totalLen = size;      
        size_t pos = 0;   //当前正在读的数据包的位置下标
        while(totalLen >= PG_CMD_PACKET_HEADLEN)  //1字节的消息类型+4字节的消息长度(包含消息长度本身)
        {
            int8_t cmdType = buf[pos];
            pos += 1;
            totalLen -= 1;
            if(cmdType == RequestType::Strartup && size > 7)
            {
                uint32_t tempLen = (uint32_t)(((buf[1] & 0xFF) << 16) | ((buf[2] & 0xFF) << 8) | (buf[3] & 0xFF));
                uint16_t tag1 = (uint16_t)((buf[4] & 0xFF) << 8) | (buf[5] & 0xFF);
                uint16_t tag2 = (uint16_t)((buf[6] & 0xFF) << 8) | (buf[7] & 0xFF);
                // SSLRequest在传输长度后传输1234与5679作为特殊标识
                // CancelRequest在传输长度后传输1234与5678作为特殊标识
                // Vertica loadBalance在传输长度后传输1234与0作为特殊标识
                if(tempLen > 2000 || !((1235 == tag1 && 0 == tag2)|| (1234 == tag1 && 5679 == tag2) || (PG_PROTOCOL_VERSION_V3 == tag1) || (1234 == tag1 && 5678 == tag2)))
                {
                    LOG_DEBUG("postgre match failed ! addNoMatchCount !")
                    addNoMatchCount();
                    return DETECT_RESULT::DETECT_FAIL;
                }
                uint64_t bodyLen = size - 4;
                if(bodyLen == PG_SSLREQUEST_BODY_LEN)
                {
                    if(tag1 == 1235 && tag2 == 0)
                    {
                        LOG_DEBUG("LoadBalanceRequest detected !");
                        m_noCmdReqStatus = NoCmdReqStatus::LoadBalance_cmd;
                    }
                    else if(1234 == tag1 && 5679 == tag2)
                    {
                        LOG_DEBUG("SSLRequest detected !");
                        m_noCmdReqStatus = NoCmdReqStatus::SSLRequest_cmd;
                    }
                    else
                    {
                        return DETECT_RESULT::DETECT_UNKNOWN;
                    }
                }
                if(bodyLen == PG_CANCELREQUEST_BODY_LEN)
                {
                    LOG_DEBUG("CancelRequest packet detected !");
                    m_noCmdReqStatus = NoCmdReqStatus::cancelRequest;
                }
                if(tag1 == PG_PROTOCOL_VERSION_V3)
                {
                    LOG_DEBUG("StartUp Packet detected !");
                    m_noCmdReqStatus = NoCmdReqStatus::Strartup_cmd;
                }
                direction = secsmart_session::session_dirction_t::SESSION_DIRCTION_REQUEST;
                //LOG_DEBUG("Startup Packet is detected!");
                if(!setIpPort(session,getDirction(),is_forward))
                {
                    LOG_DEBUG("set ip port fail");
                    return DETECT_RESULT::DETECT_FAIL;
                }
                addMatchCount();
                return DETECT_RESULT::DETECT_OK;
            }
            if(!isTrueV3Cmd(cmdType))  //检测是否是合法的postgre数据库消息类型
            {
                //这里有可能存在TCP分段的情况，属于PG的数据包却返回了false
                //因此没有调用addNoMatchCount函数，并且返回的结果是unknown,而不是fail
                return DETECT_RESULT::DETECT_UNKNOWN;
            }
            else 
            {
                uint32_t pBodyLen = (uint32_t)(((buf[pos] & 0xFF) << 24) | ((buf[pos+1] & 0xFF) << 16) | ((buf[pos+2]  & 0xFF )<< 8 ) | (buf[pos+3] & 0xFF));
                if(pBodyLen < PG_NOCMD_PACKET_HEADLEN)
                {
                    LOG_DEBUG("postgre match failed ! addNoMatchCount !")
                    addNoMatchCount();
                    return DETECT_RESULT::DETECT_FAIL;
                }
                pos += 4;     //4字节的消息长度
                totalLen -= 4;
                pBodyLen -= PG_NOCMD_PACKET_HEADLEN;  //pBodyLen 减去代表消息长度的4个字节
                    if(cmdType == RespondType::ParameterStatus)
                {
                     //提取版本信息
                     size_t tempPos = pos;
                     size_t tempLen = tempPos + pBodyLen;
                     std::string keyStr;
                     std::string valueStr;
                     for(tempPos; tempPos < tempLen && buf[tempPos] != '\0' ; tempPos++)
                     {
                        keyStr += buf[tempPos];
                     }
                     for(tempPos++; tempPos < tempLen && buf[tempPos] != '\0' ; tempPos++)
                     {
                        valueStr += buf[tempPos];
                     }
                     if(keyStr == "server_version" && !valueStr.empty())
                     {
                        version = valueStr;
                     }
                     else if(keyStr == "protocol_version" && !valueStr.empty())
                     {
                        databaseName = DBNAME_VERTICA;
                     }
                     else if(keyStr == "gp_server_version" && !valueStr.empty())
                     {
                        databaseName = DBNAME_GREENPLUM;
                     }
                     if(version.find("TBase") != std::string::npos)
                     {
                        databaseName = DBNAME_TBASE;
                     }

                }
                pos += pBodyLen;   
                totalLen -= pBodyLen;
            }

        }
        if(totalLen == 0)   //postgre消息类型解完后长度应该为0,TCP不分段的情况下
        { 
            if(!setIpPort(session,getDirction(),is_forward))
                {
                    LOG_DEBUG("set ip port fail");
                    return DETECT_RESULT::DETECT_FAIL;
                }
            //这里进行方向的判断是因为direction的值有可能为DETECT_UNKNOWN
            if(direction ==secsmart_session::session_dirction_t::SESSION_DIRCTION_REQUEST || direction == secsmart_session::session_dirction_t::SESSION_DIRCTION_RESPONSE)
            {
                addMatchCount();
                return DETECT_RESULT::DETECT_OK;
            }
            else 
            {
                return DETECT_RESULT::DETECT_UNKNOWN;
            }

        }
        else 
        {
            // TCP分段的情况不进行识别
            return DETECT_RESULT::DETECT_UNKNOWN;
        }

    }
    bool PostgreProtocol::isTrueV3Cmd(int8_t cmdType)
        {
            switch(cmdType)
            {
                case RequestType::Bind:
                case RequestType::Parse:
                case RequestType::Query:
                case RequestType::Terminate:
                case RequestType::CopyErrorRequest:
                case RequestType::CopyFail:
                case RequestType::EndOfBatchRequest:
                case RequestType::PasswordMessage:
                {
                    LOG_DEBUG("Request packet detected !")
                    direction = secsmart_session::session_dirction_t::SESSION_DIRCTION_REQUEST;
                    return true;
                }
                case RespondType::ParseComplete:
                case RespondType::BindComplete:
                case RespondType::CloseComplete:
                case RespondType::Notification:
                case RespondType::CopyIn:
                case RespondType::EmptyQuery:
                case RespondType::EndOfBatchResponse:
                case RespondType::BackendKeyData:
                case RespondType::Notice:
                case RespondType::WriteFile:
                case RespondType::Auth:
                case RespondType::RowDescription:
                case RespondType::FunctionCallResponse:
                case RespondType::CopyBoth:
                case RespondType::LoadBalance:
                case RespondType::ReadForQuery:
                case RespondType::CommandDescription:
                case RespondType::NoData:
                case RespondType::PortalSuspended:
                case RespondType::ParameterDescription:
                {
                    LOG_DEBUG("Response packet detected !")
                    direction = secsmart_session::session_dirction_t::SESSION_DIRCTION_RESPONSE;
                    return true;
                }
                // 这些case的情况是请求和响应的枚举数值一样，没法判断是请求的消息类型还是返回的消息类型
                //因此只判断是否是postgre的消息类型，对方向不做处理
                case RequestType::Close:
                case RequestType::Describe:
                case RequestType::Execute:
                case RequestType::FunctionCall:
                case RequestType::Flush:
                case RequestType::Sync:
                case RequestType::Mars:
                case RequestType::CopyDoneRequest:
                case RequestType::CopyDataRequest:
                {
                    LOG_DEBUG("PostgreSQL packet detected !,But direction cat't judge")
                    return true;
                }
                default:
                {
                    //这里有可能存在TCP分段的情况，属于PG的数据包返回了false
                    return false;
                }
                
            }
        }

    std::string PostgreProtocol::extendInfo()
    {
        return "";
    }
    
}